package leetcode.sort;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * @author wangmaolin1
 * @date 2022/6/19 6:35 下午
 */
public class _347前k个高频元素 {
    public static int[] topKFrequent(int[] nums, int k) {
        if (nums.length == 1) {
            return new int[]{nums[0]};
        }
        Map<Integer, Integer> countMap = new HashMap<>(nums.length);
        for (int num : nums) {
            if (countMap.containsKey(num)) {
                countMap.put(num, countMap.get(num) + 1);
            } else {
                countMap.put(num, 1);
            }
        }
        PriorityQueue<int[]> priorityQueue = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[1] - o2[1];
            }
        });
        for (Map.Entry<Integer, Integer> entry : countMap.entrySet()) {
            Integer num = entry.getKey();
            Integer c = entry.getValue();
            if (priorityQueue.size() == k) {
                if (priorityQueue.peek()[1] < c) {
                    priorityQueue.poll();
                    priorityQueue.offer(new int[]{num, c});
                }
            } else {
                priorityQueue.offer(new int[]{num, c});
            }
        }

        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = priorityQueue.poll()[0];
        }
        return result;
    }

    public static void main(String[] args) {
        topKFrequent(new int[]{1, 1, 1, 2, 2, 2, 3, 4, 4, 5, 2}, 2);
    }
}
